#!/usr/bin/env bash
set -euo pipefail  # 更安全的脚本执行模式

# 定义颜色代码用于输出
readonly RED='\033[0;31m'
readonly GREEN='\033[0;32m'
readonly YELLOW='\033[0;33m'
readonly NC='\033[0m' # No Color

# 配置环境变量
export HDF5_USE_FILE_LOCKING="FALSE"
export RHDF5_USE_FILE_LOCKING="FALSE"

# 加载用户环境
source ~/.bashrc

# 显示帮助信息
usage() {
    echo -e "${GREEN}Usage:${NC}
    使用pyscenic进行转录因子识别分析，需要两个参数：
    1. Seurat对象文件路径(rds格式)
    2. 输出目录
    
    结果将生成一个pyscenic_output.rds文件，包含regulons和regulonsAUC。

${YELLOW}示例:${NC}
    $0 <seurat_obj_path.rds> <output_directory>"
}

# 检查参数数量
if [ "$#" -ne 2 ]; then
    usage
    exit 1
fi

# 定义常量
readonly SEURAT_OBJ="$1"
readonly OUTPUT_DIR="$2"
readonly INPUT_LOOM="${OUTPUT_DIR}/seurat_obj_expression.loom"
readonly ADJ_TSV="${OUTPUT_DIR}/adj.sample.tsv"
readonly REG_CSV="${OUTPUT_DIR}/reg.csv"
readonly OUTPUT_LOOM="${OUTPUT_DIR}/out_SCENIC.loom"

# 定义资源文件路径 (应该通过配置文件或参数传入)
readonly TFS="/Pub/Users/wangyk/project/Poroject/F240108002_SKCM_sc/out/4.sc_houxu/十.3/scenic_data/hs_hgnc_tfs.txt"
readonly FEATHER="/Pub/Users/wangyk/project/Poroject/F240108002_SKCM_sc/out/4.sc_houxu/十.3/scenic_data/hg19-tss-centered-10kb-10species.mc9nr.genes_vs_motifs.rankings.feather"
readonly TBL="/Pub/Users/wangyk/project/Poroject/F240108002_SKCM_sc/out/4.sc_houxu/十.3/scenic_data/motifs-v9-nr.hgnc-m0.001-o0.0.tbl"

# 检查依赖工具是否存在
check_dependencies() {
    local missing=0
    
    if ! command -v pyscenic &> /dev/null; then
        echo -e "${RED}错误: pyscenic 未找到${NC}"
        missing=1
    fi
    
    if [ ! -f "/Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/SC/pysceinc/seurat2loom_exprMat.r" ]; then
        echo -e "${RED}错误: seurat2loom_exprMat.r 脚本未找到${NC}"
        missing=1
    fi
    
    if [ ! -f "/Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/SC/pysceinc/sceinc_res_loom2rds.r" ]; then
        echo -e "${RED}错误: sceinc_res_loom2rds.r 脚本未找到${NC}"
        missing=1
    fi
    
    return $missing
}

# 创建输出目录
create_output_dir() {
    if [ ! -d "${OUTPUT_DIR}" ]; then
        mkdir -p "${OUTPUT_DIR}" || {
            echo -e "${RED}无法创建输出目录: ${OUTPUT_DIR}${NC}"
            exit 1
        }
    fi
}

# 记录日志函数
log() {
    local level=$1
    local message=$2
    local timestamp=$(date +"%Y-%m-%d %H:%M:%S")
    
    case $level in
        INFO)
            echo -e "${GREEN}[${timestamp} INFO]${NC} ${message}"
            ;;
        WARN)
            echo -e "${YELLOW}[${timestamp} WARN]${NC} ${message}"
            ;;
        ERROR)
            echo -e "${RED}[${timestamp} ERROR]${NC} ${message}"
            ;;
        *)
            echo -e "[${timestamp} ${level}] ${message}"
            ;;
    esac
}

# 主函数
main() {
    log "INFO" "开始 pyscenic 分析流程"
    
    # 检查依赖
    if ! check_dependencies; then
        log "ERROR" "缺少必要的依赖"
        exit 1
    fi
    
    # 创建输出目录
    create_output_dir
    
    # 1. 转换 Seurat 对象为 loom 格式
    log "INFO" ">> 1. 转换 Seurat 对象为 loom 格式..."
    /Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/SC/pysceinc/seurat2loom_exprMat.r "${SEURAT_OBJ}" "${OUTPUT_DIR}" || {
        log "ERROR" "Seurat 对象转换失败"
        exit 1
    }
    
    # 激活 conda 环境
    source /Pub/Users/wangyk/miniconda3/bin/activate pyscenic
    
    # 2.1 运行 GRN
    log "INFO" ">> 2.1 运行 GRN (Gene Regulatory Network)..."
    pyscenic grn \
        --num_workers 36 \
        --output "${ADJ_TSV}" \
        --method grnboost2 \
        "${INPUT_LOOM}" \
        "${TFS}" || {
        log "ERROR" "GRN 步骤失败"
        exit 1
    }
    
    # 2.2 运行 cistarget
    log "INFO" ">> 2.2 运行 cistarget..."
    pyscenic ctx \
        "${ADJ_TSV}" "${FEATHER}" \
        --annotations_fname "${TBL}" \
        --expression_mtx_fname "${INPUT_LOOM}" \
        --mode "dask_multiprocessing" \
        --output "${REG_CSV}" \
        --num_workers 36 \
        --mask_dropouts || {
        log "ERROR" "cistarget 步骤失败"
        exit 1
    }
    
    # 2.3 运行 AUCell
    log "INFO" ">> 2.3 运行 AUCell..."
    pyscenic aucell \
        "${INPUT_LOOM}" "${REG_CSV}" \
        --output "${OUTPUT_LOOM}" \
        --num_workers 36 || {
        log "ERROR" "AUCell 步骤失败"
        exit 1
    }
    
    # 2.4 转换结果到 RDS 格式
    log "INFO" ">> 2.4 转换结果到 RDS 格式..."
    /Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/SC/pysceinc/sceinc_res_loom2rds.r "${OUTPUT_LOOM}" || {
        log "ERROR" "结果转换失败"
        exit 1
    }
    
    log "INFO" ">> pyscenic 分析完成!"
}

# 执行主函数
main